#!/usr/bin/python

## TODO

import os.path                  # to manipulate filenames
import sys                      # error and output messages
import re                       # checking date formats
import numpy                    # needed to hold array for output and do maths
import numpy.ma                 # needed to hold array for output and do maths
import getopt                   # CLI options

from tcclimFunctions import *   # get common functions 

def getTynSC2(model, scenario, variable, start='', end='', outbase='', datefile='', average=True, verbose=True, transposedb='', clipping=True, delta=False):
    
    # model = 'csiro2'
    # scenario = 'a2'
    # variable = 'pre'
    # start = '05,2010'
    # end = '07,2014'
    # datefile=''; average=True; verbose=True; transposedb='transpose';outbase=''
    
    # This is the demo example given in the guidance notes
    # getTynSC2('hadcm3', 'a2', 'pre', '07,2081', '07,2081', average=False)
    # test point (363, 286) should be 50.5
    
    #########################
    # SETUP AND TEST INPUTS #
    #########################
    
    # define lists of valid model options
    models    = ['cgcm2', 'csiro2', 'echam4', 'hadcm3', 'pcm']
    scenarios = ['a1','a2','b1','b2']
    variables = ['dtr','pre','tmp','vap','cld']
    
    if not model in models:
        sys.stderr.write('Model \'' + model + '\' not recognised\n')
        sys.exit()
    
    if not scenario in scenarios:
        sys.stderr.write('Scenario \'' + scenario + '\' not recognised\n')
        sys.exit()
    
    if not variable in variables:
        sys.stderr.write('Variable \'' + variable + '\' not recognised\n')
        sys.exit()
    
    if not os.path.exists(transposedb):
        sys.stderr.write('Transpose database not found at: ' + os.path.abspath(os.path.join(os.getcwd(), transposedb)) + '\n')
        sys.exit()
    
    if verbose: 
        sys.stdout.write('Obtaining dates...\n')
    
    ##########################
    # GET THE RANGE OF DATES #
    ##########################
    
    # compile a regular expression to check for badly formatted count fields
    reMoYr = re.compile('[0-9]{1,2},[0-9]{4}') 
    
    # turn start and end into a list or read from file
    if datefile != '':
        # try and get the file contents
        try:
            dateFile = open(datefile, 'r')
        except IOError:
            sys.stderr.write('Cannot open file containing dates: ' + datefile + '\n')
            sys.exit()
        
        dateKeys = dateFile.readlines()
        dateKeys = [m.strip() for m in dateKeys]
        dateFile.close()
        
        # test for carriage return at end
        if dateKeys[-1] == '':
            dateKeys = dateKeys[:-1]
        
        # test for two parts
        dates = [m.split(',') for m in dateKeys]
        nPart = [len(m) == 2 for m in dates]
        if not all(nPart):
            sys.stderr.write('Badly formatted dates in: ' + datefile + '\n')
            sys.exit()
        
        # try and convert pairs of values to integers
        try:
            dates = [(int(m[0]), int(m[1])) for m in dates]
        except:
            sys.stderr.write('Datefile contains non-integer values: ' + datefile + '\n')
            sys.exit()
        
        # check their range
        monthValid = [m[0] > 0 and m[0] <= 12  for m in dates]
        yearValid = [m[1] > 2000 and m[1] <= 2100  for m in dates]
        
        if not all(monthValid):
            sys.stderr.write('Month in ' + datefile + ' outside range [1-12].\n')
            sys.exit()
        if not all(yearValid):
            sys.stderr.write('Year in ' + datefile + ' outside range [2001,2100].\n')
            sys.exit()
        
        # remove duplicates
        noDupes = list(set(dates))
        if len(noDupes) < len(dates):
            sys.stderr.write('WARNING: duplicate dates removed from dates in ' + datefile + '\n')
            dates = noDupes
            dates.sort()
    else: 
        # are both start and end defined and defined correctly
        if start == '' or end == '':
            sys.stderr.write('Missing start or end date.\n')
            sys.exit()
        
        if not reMoYr.match(start):
            sys.stderr.write('Misformatted start date: ' + start + '\n')
            sys.exit()
        
        if not reMoYr.match(end):
            sys.stderr.write('Misformatted end date: ' + end + '\n')
            sys.exit()
        
        # check the dates make sense
        start = map(int, start.split(','))
        end   = map(int, end.split(','))
        if start[0] > 12:
            sys.stderr.write('Bad start month: ' + str(start[0]) + '\n')
            sys.exit()
        
        if end[0] > 12:
            sys.stderr.write('Bad end month: ' + str(end[0]) + '\n')
            sys.exit()
        
        if start[1] < 2001 or start[1] > 2100:
            sys.stderr.write('Start year not in range 2001 - 2100: ' + str(start[1]) + '\n')
            sys.exit()
        
        if end[1] < 2001 or end[1] > 2100:
            sys.stderr.write('End year not in range 2001 - 2100: ' + str(end[1]) + '\n')
            sys.exit()
        
        # get the date range
        if (start[1] > end[1]) or (start[1] == end[1] and start[0] > end[0]):
            sys.stderr.write('Start date later than end date\n')
            sys.exit()
        
        start_month = start[0]
        nMonths     = (end[1] - start[1]) * 12 + end[0] + 1
        dates = [((m - 1) % 12 + 1, (m - 1) / 12 + start[1]) for m in range(start_month, nMonths)]
    
    # reconsitute dateKeys from parsed dates - has the effect of stripping leading zeros
    dateKeys = [str(m[0]) + ',' + str(m[1]) for m in dates]
    
    ########################
    # READ DATA FROM FILES #
    ########################
    
    # dictionary of base HD file names - some small differences in model choice
    hdFileDict = {'cgcm2':  {'a1': 'cgcm2.a1fi.2080s', 'a2': 'cgcm2.a2a.2080s',  'b1': 'cgcm2.b1.2080s',  'b2': 'cgcm2.b2a.2080s' },
                  'csiro2': {'a1': 'csiro2.a1a.2080s', 'a2': 'csiro2.a2a.2080s', 'b1': 'csiro2.b1a.2080s','b2': 'csiro2.b2a.2080s'},
                  'echam4': {'a1': '',                 'a2': 'echam4.a2.lan',    'b1': '',                'b2': 'echam4.b2.lan'   },
                  'hadcm3': {'a1': 'had3.a1fi.2080s',  'a2': 'had3.a2m.2080s',   'b1': 'had3.b1.2080s',   'b2': 'had3.b2m.2080s'  },
                  'pcm':    {'a1': 'pcm.a1fi.2080s',   'a2': 'pcm.a2a.2080s',    'b1': 'pcm.b1.2080s',    'b2': 'pcm.b2a.2080s'   }}
    
    # dictionary of base ann file names
    annScenarios = {'a1':'a1fi','a2':'a2','b1':'b1','b2':'b2'}
    
    # read translated hd file 
    if verbose: sys.stdout.write('Reading 2071-2099 projections (hd)...\n')
    
    hdFile = os.path.join(transposedb, 'tr.hd.' + hdFileDict[model][scenario] + '.' + variable)
    hd = readTransposeGrim(hdFile)
    
    # read ann file
    # - time series of global temperature by year
    if verbose: sys.stdout.write('Reading temperature time series (ann)...\n')
    
    annFile = os.path.join(transposedb, model + '-' + annScenarios[scenario] + '.ann')
    ann = readAnn(annFile)
    
    # - read obs file unless running in delta mode:
    # - monthly maps of observed values (1961-1990) for variables
    # - dated [month],1990
    if delta:
        if verbose: sys.stdout.write('Skipping reading observed values(obs) and writing a delta file...\n')
    else:
        if verbose: sys.stdout.write('Reading 1961-1990 observed values (obs)...\n')
        
        obsFile = os.path.join(transposedb, 'tr.obs.clim6190.globe.lan.' + variable)
        obs = readTransposeGrim(obsFile)
    
    # get keys for iavar years - future - 100
    iaDateKeys = [str(m[0]) + ',' + str(m[1]- 100) for m in dates]
    
    # read iavar file:
    # - this is time consuming because the files are so large
    #   but this now only returns the dates needed
    if verbose: sys.stdout.write('Reading variation values (iavar) for ' + str(len(iaDateKeys)) + ' months...\n')
    
    iavarFile = os.path.join(transposedb, 'tr.iavar.1901-2000.' + variable)
    iavar = readTransposeGrim(iavarFile, iaDateKeys, ticker=verbose)
    
    ############################
    # DO COMPUTATION OF VALUES #
    ############################
    
    if verbose: sys.stdout.write('Starting calculations of grid values for ' + str(len(iaDateKeys)) + ' months\n')
    
    # - dictionary of average temperature changes in 2071 - 2099.
    #   matches the same from ann files except echam4  and cgcm2-a1 = 4.384
    modDict = {'hadcm3': {'a2': 3.931, 'a1': 4.863, 'b2': 3.070, 'b1': 2.521}, 
               'csiro2': {'a2': 3.938, 'a1': 4.855, 'b2': 3.139, 'b1': 2.592}, 
               'echam4': {'a2': 1.000, 'a1': 1.000, 'b2': 1.000, 'b1': 1.000},
               'cgcm2':  {'a2': 3.548, 'a1': 4.382, 'b2': 2.462, 'b1': 2.023},
               'pcm':    {'a2': 2.462, 'a1': 3.045, 'b2': 1.894, 'b1': 1.541}}
    
    # 1) normalise HD to get average variable change per degree C
    hd.update((k, v / modDict[model][scenario]) for k, v in hd.iteritems())
    
    # now process individual dates:
    
    # store data in a numpy array
    # - this _does_ create 2d array with only 1 date
    valArray = numpy.zeros([len(dates), 67420])
    
    for currDate in range(len(dates)):
        
        if verbose:            
            sys.stdout.write('.')
            if ((currDate+1) % 40) == 0: sys.stdout.write('\n')
            sys.stdout.flush()
        
        # get all the indices together
        # this assumes that dates, dateKeys and iaDateKeys are in the same order
        c_month = dates[currDate][0]
        c_year  = dates[currDate][1]
        c_ia    = iaDateKeys[currDate]
        c_obs   = str(c_month) + ',1990'
        c_hd    = str(c_month) + ',2080'
        c_key   = dateKeys[currDate]
                
        # 2) cross multiply normalised hd for this month (hd[c_hd])
        #    by temperature for this year (ann[c_year])
        #    to get projected global change for that year
        c_data = hd[c_hd] * ann[c_year]
        
        # 3) add inter-annual and decadal variation from iavar
        #    to each grid ref in parallel
        c_data = c_data + iavar[c_ia]
        
        # 4) add observed mean value for 1961-90 reference and put into array
        if delta:
            valArray[currDate,] = c_data
        else:
            valArray[currDate,] = c_data + obs[c_obs]
    
    if verbose: sys.stdout.write('\n')
    
    # 5) Scale all the data - all of the Tyn SC 2.03 variables are scaled by 0.1
    valArray = valArray * 0.1
    
    # 6) handle clipping of all out of range values
    #    - need to use masked array values for this
    #      to get sensible averages in the presence of invalid data
    if verbose: sys.stdout.write('Handling out of range data...\n')
    
    # - dictionaries of min and max values
    minVar = {'tmp': -273.15, 'pre': 0, 'dtr': 0.1, 'vap': 0, 'cld': 0}
    maxVar = {'tmp': numpy.Inf, 'pre': numpy.Inf, 'dtr': numpy.Inf, 'vap': numpy.Inf, 'cld': 100}
    
    # don't clip on delta files
    if not delta:
        if not clipping:
            # blank out out of range values
            valArray = numpy.ma.masked_greater(valArray, maxVar[variable])
            valArray = numpy.ma.masked_less(valArray, minVar[variable])
        
            # warn about this if encountered
            if valArray.mask.any() == True:
                warn = '\n  WARNING:\n  Modelled outputs include out of range values.\n  ' + \
                          str(sum(sum(valArray.mask))) + ' out of ' + str(valArray.size) + ' calculated values have been treated as missing\n  ' + \
                         'These values will be reported as missing data in ascii grid files\n  ' + \
                         'and will not be included in calculation of average grids.\n\n  ' + \
                         'This is common for some models and variables.\n\n' 
                sys.stderr.write(warn)
        else:
            # clip out of range values
            valArrayTooHigh = valArray > maxVar[variable]
            valArray[valArrayTooHigh] = maxVar[variable]
            valArrayTooLow = valArray < minVar[variable]
            valArray[valArrayTooLow] = minVar[variable]
        
            # set a mask with no masking
            valArray = numpy.ma.masked_array(valArray)
        
            # warn about this
            if valArrayTooHigh.any() or valArrayTooLow.any():
                warn = '\n  WARNING:\n  Modelled outputs include out of range values.\n  ' + \
                         str(sum(sum(valArrayTooLow)) + sum(sum(valArrayTooHigh))) + ' calculated values have been clipped at the minimum or maximum value.\n  ' + \
                         'This is common for some models and variables.\n\n' 
                sys.stderr.write(warn)
    
    # average across the dates if required (axis 0 = columns)
    #  - returns a masked value if _all_ values in a column are masked
    #    but otherwise returns a non-masked array, so need to remask
    #  - TODO add a map of N to show where the missing data are
    if average:
        if verbose: sys.stdout.write('Averaging values across dates...\n')
        valArray = numpy.average(valArray, axis=0)
        valArray = numpy.ma.masked_array(valArray)
        
    
    # handle singleton year extractions and averages, where numpy
    # drops dimensions, by putting the dimension back in
    if len(valArray.shape) == 1:
        valArray = valArray[numpy.newaxis, :]
    
    ##########################
    # WRITE OUTPUTS TO FILES #
    ##########################
    
    # get the base filename
    if outbase == '':
        if delta:
            outbase = model + '.' + scenario  + '.' + variable  + '.delta.'
        else:
            outbase = model + '.' + scenario  + '.' + variable  + '.'
    
    # now output the data as ascii files
    
    # reusable bits
    header = "ncols        720\nnrows        360\nxllcorner    -180\nyllcorner    -90\ncellsize     0.5\nNODATA_value -999\n"
    rowOrder = range(360)
    rowOrder.reverse()
    
    # get a set of dates to put in filename
    if average:
        if datefile == '':
            dateKeys = ['ave-' + str(dates[0][0]) + '.' + str(dates[0][1]) + '-' + str(dates[-1][0]) + '.' + str(dates[-1][1]) ]
        else:
            dateKeys = ['ave-' + os.path.basename(datefile)]
    else:
        # replace commas in dateKeys by '.'
        dateKeys = [str(m[0]) + '.' + str(m[1]) for m in dates]
    
    # get the grid indices from gridindex
    gridfile= transposedb + '/' + 'transposedGridRefIndex.txt'
    try:
        gridind = open(gridfile, 'r')
    except IOError:                                                  
        sys.stderr.write('Cannot open grid index file: ' + gridindex + '\n')
        sys.exit()
    
    gridindex = gridind.readlines()
    gridind.close()
    gridindex = [ x.rstrip() for x in gridindex]
    gridindex = [ x.split(',') for x in gridindex]
    gridindex = [ [int(row), int(col)]  for row, col in gridindex]
    
    ## now that we are writing to file, may need to turn masked values
    ## in valArray into missing data values
    if valArray.mask.any() == True:
        valArray.fill_value = -999
        valArray = valArray.filled()
    
    if verbose: sys.stdout.write('Writing output files for '+ str(len(dateKeys)) + ' ascii files\n')
    
    # iterate over dates    
    for eachkey in range(len(dateKeys)):
        
        # generate a grid prefilled with -999
        ascii = numpy.zeros([360,720]) - 999
        
        # fill in asciiList
        for c in range(len(gridindex)):
            row = gridindex[c][1] - 1
            col = gridindex[c][0] - 1
            ascii[row , col] = valArray[eachkey, c]
        
        # write out asciiList as a grid
        outfile = outbase + dateKeys[eachkey] + '.asc'
        
        try:
            out = open(outfile, 'w')
        except IOError:                                                  
            sys.stderr.write('Cannot open output file: ' + outfile + '\n')
            sys.exit()
        
        # write header to file
        out.write(header)
        
        # write rows to file
        for r in rowOrder:
            # space delimited strings
            rowVals = [' '+ str(round(v, 2)) for v in ascii[r,]]
            rowVals = ''.join(rowVals)
            # lose first space 
            rowVals = rowVals[1:]
            out.write(rowVals + '\n')
        
        out.close()
        
        # file ticker
        if verbose: 
            sys.stdout.write('.')
            sys.stdout.flush()
    
    # end ticker
    sys.stdout.write('\n')


def usage():
    
    print """
  
  tcclimTynSC2:
  A data extraction tool for Tyndall SC 2.03 global climate projections.
  
  Basic use:
     tcclimTynSC2.py [options] 'model' 'scenario' 'variable' 'startdate' 'enddate'
  
  Options:
     -h or --help:        show this help information
     -q or --quiet:       suppress normal progress reporting
     -m or --monthly:     provide monthly grid files not a single average file    
     -d or --datefile:    use dates from provided datefile
     -o or --outbase:     use provided base name for output files
     -t or --transposedb: use provided path to transposed data folder 
     -c or --clipping:    truncate out of range projection values rather 
                          than returning missing values
     -D or --delta:       return projected anomaly from baseline not projected value 
  
  Available models:
     'cgcm2', 'csiro2', 'echam4', 'hadcm3', 'pcm'
  Available scenarios:
     'a1','a2','b1','b2'
  Available variables:
     'dtr','pre','tmp','vap','cld'
  Date format and range:
     '01,2001' to '12,2100'
  
  Returns ArcGIS ascii grid files of global data for a given model,
  scenario and variable. The user can provide  a range of dates
  from the command line. Alternatively, the user can name a file containing
  dates, with each line holding a single date in the format above. Dates
  provided in this way need not be consecutive.
  
  By default, the program calculates an average across the dates selected
  but monthly files may be requested instead. As each output file is around
  2 MB, this can be slow and use a lot of disk space for large date ranges.
  
  The climate projections sometimes fall outside of sensible ranges for the
  variables, for example, a precipitation of less than zero or cloud cover of
  more than 100%. The default action is for such values to be converted to 
  missing data, but the -c option sets these values to be clipped at the 
  relevant minimum or maximum value.
  
  The delta option returns only the projected change from the baseline values 
  from 1961-1990. It includes the interannual variation in the predicted anomaly
  but does not add the anomaly to the baseline data. This mode does not identify
  anomalies  that lead to nonsensical absolute prediction: no values are 
  clipped or converted to missing data.
  
  The tool requires a database of transposed grim files created using the 
  script 'tcclimDownloadTYNSC2.py'. This transposition radically improves 
  handling speed for converting the data to ascii grids. By default, the tool
  looks for a folder called 'tcclim_tyn_sc_2_03' in the directory from 
  which the tool is being run.
  
  Example use:
     # Get precipitation from HadCM3 A2 for July 2081
     # This is the worked example from:
     # http://www.cru.uea.ac.uk/cru/data/hrg/timm/grid/algorithm.html
     # Grid cell: row 286, col 363 from lower left origin = 50.5
     tcclimTynSC2.py -m hadcm3 a2  pre  07,2081  07,2081
    
     # Get mean precipitation from HadCM3 A2 from July 2071 to July 2091
     tcclimTynSC2.py hadcm3 a2  pre  07,2071  07,2091
     
     # Get monthly grid files for HadCM3 A2 for dates in a file
     tcclimTynSC2.py -m -d exampleDatesTynSC2.txt hadcm3 a2  pre
    
    """


def main():
    
    # Tidies up and passes options to this function:
    # def getTynSC2(model, scenario, variable, start='', end='', 
    #               outbase='', datefile='', average=True, verbose=True, 
    #               transposedb='tcclim_tyn_sc_2_03'):
    
    # complex input to test control
    # python getTynSC2.py -mq -d dates.txt -o 'myFiles' --transposedb '../transpose' 'hadcm3' 'a2'  'pre'
    
    # sort out the options and argument
    try:
        opts, args = getopt.getopt(sys.argv[1:], "hqmcDd:o:t:", ["help", "quiet","monthly","clipping","delta","datefile=","outbase=", "transposedb="])
    except getopt.GetoptError, err:
        # print help information and exit:
        print str(err) # will print something like "option -a not recognized"
        usage()
        sys.exit()
    
    # default options
    verbose=True
    average=True
    datefile=''
    outbase =''
    transposedb='tcclim_tyn_sc_2_03'
    clipping = False
    delta=False
    
    # initialise blank dates for handling datefile option
    start=''
    end=''
    
    # initial setups
    for o, a in opts:
        if o in ('-h', '--help'):
            usage()
            sys.exit()
        elif o in ('-q', '--quiet'):
            verbose=False
        elif o in ('-m', '--monthly'):
            average=False
        elif o in ('-D', '--delta'):
            delta=True
        # getopt handles missing arguments for d, o and t but can easily pass model
        # and scenario names in as datefiles. Let getTynSC2 function handle this
        elif o in ('-d', '--datefile'):
            datefile = a
        elif o in ('-o', '--outbase'):
            outbase = a
        elif o in ('-t', '--transposedb'):
            transposedb = a
        elif o in ('-c', '--clipping'):
            clipping=True
        else:
            assert False, "Unhandled option"
            usage()
            sys.exit()
    
    if len(args) == 3:
        mod = args[0]
        sce = args[1]
        var = args[2]
    elif len(args) == 5:
        mod = args[0]
        sce = args[1]
        var = args[2]
        start = args[3]
        end = args[4]
    else:
        sys.stderr.write('\n  Incorrect number of arguments:\n')
        sys.stderr.write('  getTynSC2 require either three arguments (model, scenario, variable) \n')
        sys.stderr.write('  and a datefile or five arguments (model, scenario, variable,\n')
        sys.stderr.write('  start date, end date)\n')
        usage()
        sys.exit()
    
    ## [sys.stdout.write(o + ',' + a + '\n') for o, a in opts]
    ## [sys.stdout.write(a + '\n') for a in args]
    
    getTynSC2(model=mod, scenario=sce, variable=var, start=start, end=end,  
              outbase=outbase, datefile=datefile, average=average, verbose=verbose,
              transposedb=transposedb, clipping=clipping, delta=delta)


if __name__ == "__main__":
    main()
